package record

import (
	"errors"
	"github.com/fcbox/dbreplay/dbreplay-server/replay"
	"github.com/fcbox/dbreplay/dbreplay-server/util"
	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"
	"github.com/google/gopacket/pcap"
	"io"
	"log"
	"strconv"
	"strings"
	"time"
)

var(
	byteMap map[string] *[]byte = make(map[string] *[]byte)
)
type RecordConf struct {
	Ip  string
	Port int
	Tip string
	Tport int
	Tuser string
	Tpassword string
	Tdb string
	Speed int
	Replay int
	FileObj io.Writer
	ErrFileObj io.Writer
	replayObj *replay.Replay
}

func NewRecordConf() *RecordConf {
	record := new(RecordConf)
	return record
}

func(recon * RecordConf) Watch() error  {
	var device string = ""
	var snapshot_len   int32  = 65535
	var err            error
	var timeout        time.Duration = 30 * time.Second
	var handle         *pcap.Handle
	devices, err := pcap.FindAllDevs()
	if err != nil {
		log.Fatal(err)
	}
	for _, deviceGet := range devices {
		for _, address := range deviceGet.Addresses {
			if address.IP.String() == recon.Ip {
				device = deviceGet.Name
				break
			}
		}
	}
	if device == "" {
		err = errors.New("the ip don't match the network device. maybe you don't have the sudo authority.")
		return err
	}
	// Open device
	handle, err = pcap.OpenLive(device, snapshot_len, true, timeout)
	if err != nil {
		log.Fatal(err)
		return err
	}
	defer handle.Close()

	var filter string = "tcp and dst host "
	filter = filter + recon.Ip + " and dst port " + strconv.Itoa(recon.Port)

	//var filter string = "tcp and host "
	//filter = filter + recon.Ip + " and port " + strconv.Itoa(recon.Port)
	err = handle.SetBPFFilter(filter)
	if err != nil {
		log.Fatal(err)
	}
	if recon.Replay == 1 {
		replay := replay.NewReplay()
		replay.Tdb = recon.Tdb
		replay.Tport = recon.Tport
		replay.Tip = recon.Tip
		replay.Tpassword = recon.Tpassword
		replay.Tuser = recon.Tuser
		replay.FileObj = recon.FileObj
		replay.ErrFileObj = recon.ErrFileObj
		recon.replayObj = replay
		replay.Init()
	}
	packetSource := gopacket.NewPacketSource(handle, handle.LinkType())
	var count int = 0
	for packet := range packetSource.Packets() {
		count++
		//time.Sleep(1*time.Second)
		recon.dispatch(packet)
	}
	return err
}

func(recon * RecordConf) dispatch(packet gopacket.Packet)  {
	ipLayer := packet.Layer(layers.LayerTypeIPv4)
	tcpLayer := packet.Layer(layers.LayerTypeTCP)
	applicationLayer := packet.ApplicationLayer()
	if applicationLayer != nil && tcpLayer != nil && ipLayer != nil {
		ip, _ := ipLayer.(*layers.IPv4)
		tcp,_ :=tcpLayer.(*layers.TCP)
		/*
		if tcp.SrcPort == layers.TCPPort(recon.Tport) {
			util.WriteString(recon.FileObj,"here is output.")
			util.WriteString(recon.FileObj,util.BytesToString(applicationLayer.Payload())+"\n")
			return
		}

		 */
		payload := applicationLayer.Payload()
		originB := byteMap[ip.SrcIP.String()+":"+tcp.SrcPort.String()]
		if originB != nil {
			payload = append(*originB,payload...)
			byteMap[ip.SrcIP.String()+":"+tcp.SrcPort.String()] = nil
		}
		special := util.GetCommand(payload)

		if strings.EqualFold(special,"NEEDMORE"){
			//util.WriteString(recon.FileObj,"NEEDMORE")
			byteMap[ip.SrcIP.String()+":"+tcp.SrcPort.String()] = &payload
			return
		}
		if strings.EqualFold(special,"ERROR"){
			//util.WriteString(fileObj,util.BytesToString(applicationLayer.Payload()))
			util.WriteString(recon.ErrFileObj,util.BytesToString(applicationLayer.Payload()))
			util.WriteString(recon.ErrFileObj,"session:")
			util.WriteString(recon.ErrFileObj,ip.SrcIP.String()+":"+tcp.SrcPort.String()+":"+strconv.FormatInt(time.Now().UnixNano() / 1e6,10)+"\n")
			util.WriteString(recon.ErrFileObj,special+"\n")
			return
		}
		if recon.Replay != 1 {
			util.WriteString(recon.FileObj,"session:")
			util.WriteString(recon.FileObj,ip.SrcIP.String()+":"+tcp.SrcPort.String()+":"+strconv.FormatInt(time.Now().UnixNano() / 1e6,10)+"\n")
			util.WriteString(recon.FileObj,special+"\n")
		}else {
			//util.WriteString(recon.FileObj,"session:")
			//util.WriteString(recon.FileObj,ip.SrcIP.String()+":"+tcp.SrcPort.String()+":"+strconv.FormatInt(time.Now().UnixNano() / 1e6,10)+"\n")
			//util.WriteString(recon.FileObj,special+"\n")

			recon.replayObj.Dispatch(ip.SrcIP.String()+":"+tcp.SrcPort.String(),special)
		}
	}
	return
}